////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//                                                                                                                            //
//                                                         WNProject                                                          //
//                                                                                                                            //
//            This file is distributed under the BSD 2-Clause open source license. See License.txt for details.               //
//                                                                                                                            //
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#include "WNCore/inc/WNBase.h"
#include "WNMemory/inc/WNMemory.h"
#include "WNStrings/inc/WNStrings.h"
#include "WNScripting/inc/WNFunctionManager.h"
#include "WNScripting/inc/WNIdentifierMap.h"
#include "WNScripting/inc/WNLLVMIncludes.h"
#include "WNScripting/inc/WNScriptingEngineData.h"
#include "WNScripting/inc/WNTypeManager.h"

#ifdef _WN_WINDOWS
    #include <windows.h>
#endif

//Sadly cannot specify any user-data with our LazyFunctionCreator.. 
//only options for specifying our WNFunctionManager is global or threadlocal
//neither of which I really like. Since ideally I would like to have the ability
//to run as many scripting engines as possible.
//I am guessing for now that the option will be "1 scripting engine / thread"
__WN_THREAD_LOCAL WNFunctionManager* tlFunctionManager;
__WN_THREAD_LOCAL WN_VOID* tlTempStackLocation;
__WN_THREAD_LOCAL WN_SIZE_T tlTempStackSize;

typedef WN_VOID*(*LazyFunctionCreator)(const std::string&);

WN_VOID* GetTempStackLocation(WN_SIZE_T _size) {
    if(tlTempStackSize < _size) {
        tlTempStackSize = _size;
        tlTempStackLocation = WNRealloc(tlTempStackLocation, _size);
    }
    return(tlTempStackLocation);
}

WN_VOID* WNFunctionManager::CheckSystemFunctions(const std::string& _function) {
    #ifdef _WN_WINDOWS
        #ifdef _WN_X86
            #ifdef _WN_32_BIT
                if (_function == "_chkstk") {
                    HMODULE m = GetModuleHandle("ntdll.dll");

                    return(GetProcAddress(m, "_chkstk"));
                }
            #elif defined _WN_64_BIT
                if (_function == "__chkstk") {
                    HMODULE m = GetModuleHandle("ntdll.dll");

                    return(GetProcAddress(m, "__chkstk"));
                }
            #endif
        #endif
    #endif

    if (_function == "++GenerateNewFunction") {
        return(reinterpret_cast<WN_VOID*>(&StaticGenerationFunction));
    } else if (_function == "++malloc") {
        return(reinterpret_cast<WN_VOID*>(mInternalMallocFunc));
    } else if (_function == "++free") {
        return(reinterpret_cast<WN_VOID*>(mInternalFreeFunc));
    } else if (_function == "++alignedMalloc") {
        return(reinterpret_cast<WN_VOID*>(mInternalAlignedMalloc));
    } else if (_function == "++alignedFree") {
        return(reinterpret_cast<WN_VOID*>(mInternalAlignedFree));
    } else if (_function == "memcpy") {
        return(reinterpret_cast<WN_VOID*>(mInternalMemcpy));
    } else if (_function == "++tempStack") {
        return(reinterpret_cast<WN_VOID*>(&GetTempStackLocation));
    }

    return(WN_NULL);
}

WN_VOID WNFunctionManager::SetFreeFunction(WNScriptFreeFunc _freeFunc) {
    mInternalFreeFunc = _freeFunc;
}

WN_VOID WNFunctionManager::SetMallocFunction(WNScriptAllocFunc _scriptFunc) {
    mInternalMallocFunc = _scriptFunc;
}

WN_VOID WNFunctionManager::SetAlignedMallocFunction(WNScriptAlignedMallocFunc _mallocFunc) {
    mInternalAlignedMalloc = _mallocFunc;
}

WN_VOID WNFunctionManager::SetAlignedFreeFunction(WNScriptAlignedFree _freeFunc) {
    mInternalAlignedFree = _freeFunc;
}

WN_VOID WNFunctionManager::SetMemcpyFunction(WNScriptMemcpy _memcpy) {
    mInternalMemcpy = _memcpy;
}

WN_VOID* WNLazyFunctionCreator(const std::string& _function) {
    WNFunctionManager* mgr = tlFunctionManager;
    WN_VOID * sysFun = mgr->CheckSystemFunctions(_function);
    if (sysFun) {
        return(sysFun);
    }
    
    const WN_CHAR* separator = WNStrStr(_function.c_str(), "@");
    WN_SIZE_T len = separator - _function.c_str();
    WN_BOOL parameterized = WN_TRUE;
    if (!separator) {
        parameterized = WN_FALSE;
        len = WNStrLen(_function.c_str());
    } else if (*(separator + 1) == '\0') {
        parameterized = WN_FALSE;
    }
    WNIdentifierElement ident = mgr->GetIdentifierMap()->GetIdentifier(_function.c_str(), len);
    WN_SIZE_T numParams = 0;
    if(parameterized) {
        separator++;
        WN_SIZE_T numArray = 0;
        while(separator[numParams] != '\0') {
            if(separator[numParams] == '$') {
                numArray += 1;
            }
            numParams++;
        }
        numParams -= numArray;
    }

    std::vector<WNFunction*> possibles;
    mgr->GetExternalFunctions(ident, possibles, WN_TRUE);
    std::vector<WNFunction*>::iterator i = possibles.begin();
    
    if (parameterized) {
        for(; i != possibles.end(); ++i) {
            if ((*i)->mTypes.size() != numParams) {
                continue;
            }
            WN_BOOL found = WN_TRUE;
            for(WN_UINT32 j = 0, k = 0; j < numParams; ++j, ++k) {
                WN_BOOL isArray = separator[k] == '$';
                if(isArray) {
                    ++k;
                }
                //If our identifier is different, or we are not the same array type
                if ((*i)->mTypes[j]->mCharacterIdentifier != separator[k] ||
                    (isArray != (*i)->mTypes[j]->mArray) ){
                    found = WN_FALSE;
                    break;
                }
            }
            if(found) {
                break;
            }
        }
        if (i == possibles.end()) {
            return(WN_NULL);
        }
    }
    
    return((*i)->mFunctionPointer);
}

WNFunctionManager::WNFunctionManager(WNScriptingEngineData& _data):
                                                mScriptingData(_data) {
    mScriptingData.mExecutionEngine->DisableSymbolSearching(WN_TRUE);
    mScriptingData.mExecutionEngine->InstallLazyFunctionCreator(&WNLazyFunctionCreator);
    AddNewThread();

    llvm::Type* ptrType = llvm::PointerType::getUnqual(_data.mTypeManager->GetIntPtrType());
    llvm::Type* intType = _data.mTypeManager->GetIntPtrType();
    std::vector<llvm::Type*> llvmTypes;
    llvmTypes.push_back(ptrType);
    llvmTypes.push_back(ptrType);
    CreateInternalFunction(mCreateFunctionfunction, "++GenerateNewFunction", reinterpret_cast<WN_VOID*>(&StaticGenerationFunction), ptrType, llvmTypes);

    llvmTypes.clear();
    llvmTypes.push_back(_data.mTypeManager->GetIntPtrType());
    CreateInternalFunction(mMallocFunction, "++malloc", reinterpret_cast<WN_VOID*>(&WNMalloc), ptrType, llvmTypes);
    llvmTypes.push_back(_data.mTypeManager->GetIntPtrType());
    CreateInternalFunction(mAlignedMallocFunction, "++alignedmalloc", reinterpret_cast<WN_VOID*>(&WNAlignedMalloc), ptrType, llvmTypes);

    llvmTypes.clear();
    llvmTypes.push_back(ptrType);
    CreateInternalFunction(mFreeFunction, "++free", reinterpret_cast<WN_VOID*>(&WNFree), llvm::Type::getVoidTy(llvm::getGlobalContext()), llvmTypes);
    CreateInternalFunction(mAlignedFreeFunction, "++alignedfree", reinterpret_cast<WN_VOID*>(&WNAlignedFree), llvm::Type::getVoidTy(llvm::getGlobalContext()), llvmTypes);

    llvmTypes.clear();
    llvmTypes.push_back(ptrType);
    llvmTypes.push_back(ptrType);
    llvmTypes.push_back(_data.mTypeManager->GetIntPtrType());
    CreateInternalFunction(mMemcpyFunction, "++memcpy", reinterpret_cast<WN_VOID*>(&memcpy), ptrType, llvmTypes);

    llvmTypes.clear();
    llvmTypes.push_back(_data.mTypeManager->GetIntPtrType());
    CreateInternalFunction(mInternalTempStackFunction, "++tempStack", reinterpret_cast<WN_VOID*>(&GetTempStackLocation), ptrType, llvmTypes);

    mInternalAlignedMalloc = &WNAlignedMalloc;
    mInternalAlignedFree  = &WNAlignedFree;
    mInternalFreeFunc     = &WNFree;
    mInternalMallocFunc   = &WNMalloc;
    mInternalMemcpy       = &memcpy;
}

WNFunctionManager::~WNFunctionManager() {
    if(mCreateFunctionfunction) {
        WN_DELETE(mCreateFunctionfunction);
    }  

    for(FunctionMMap::iterator i = mNativeMap.begin(); i != mNativeMap.end(); ++i) {
        WN_DELETE((*i).second);
    }

    for(FunctionMMap::iterator i = mFunctionMap.begin(); i != mFunctionMap.end(); ++i) {
        WN_DELETE(i->second);
    }
}

WN_VOID WNFunctionManager::AddIncludeFile(WNIdentifierElement _file, WNIdentifierElement _includedFile) {
    mIncludeMap.insert(std::make_pair(_file, _includedFile));
}

WNFunctionManager::ReturnCode WNFunctionManager::AddExternalFunction(WN_VOID* _funcPtr, WN_BOOL _stack, const WN_CHAR* _returnType, const WN_CHAR* _name, const WN_CHAR* _arglist) {
    std::vector<WNNativeElementType*> functionTypes;
    const WN_CHAR* startArg = WN_NULL;
    std::string types = "@";

    for(const WN_CHAR* ch = _arglist;; ++ch) {
        if(startArg == WN_NULL) {
            if (*ch != ',' && *ch != '\n' && *ch != '\t' && *ch != ' ') {
                startArg = ch;
            }
        }
        else if (*ch == ',' || *ch == '\n' || *ch == '\t' || *ch == ' ' || *ch == '\0') {
            if ((ch - startArg)< 2) {
                return(WN_INVALID_IDENTIFIER);
            }
            WN_BOOL isArray = WN_FALSE;
            WNIdentifierElement ident;
            WNNativeElementType* nativeType;
            if(WNMemCmp(startArg, "Array<", 6) == 0) {
                isArray = WN_TRUE;
                ident = mScriptingData.mIdentifierMap->GetIdentifier(startArg + 6, (ch - startArg - 7));
                nativeType = mScriptingData.mTypeManager->GetNativeArrayFromWNType(ident);
            } else {   
                ident = mScriptingData.mIdentifierMap->GetIdentifier(startArg, (ch - startArg));
                nativeType = mScriptingData.mTypeManager->GetNativeFromWNType(ident);
            }
             
            
            if (!nativeType) {
                return(WN_INVALID_TYPE);
            }
            if(isArray) {
                types.push_back('$');
            }
            types.push_back(nativeType->mCharacterIdentifier);
            functionTypes.push_back(nativeType);
            startArg = WN_NULL;
        }
        if (*ch == '\0') {
            break;
        }
    }
    std::string strName = _name;

    WNIdentifierElement functionName = mScriptingData.mIdentifierMap->GetIdentifier(strName.c_str(), strName.size());
    strName += types;
    WNIdentifierElement correctFunctionName = mScriptingData.mIdentifierMap->GetIdentifier(strName.c_str(), strName.size());
    if(!functionName) {
        return(WN_INVALID_IDENTIFIER);
    }

    WNIdentifierElement returnTypeIdent;
    WNNativeElementType* nativeReturn;
    bool isArray = WN_FALSE;
    if(WNMemCmp(_returnType, "Array<", 6) == 0) {
        isArray = WN_TRUE;
        returnTypeIdent = mScriptingData.mIdentifierMap->GetIdentifier(_returnType + 6, WNStrLen(_returnType) - 7);
        nativeReturn = mScriptingData.mTypeManager->GetNativeArrayFromWNType(returnTypeIdent);
    } else {
        returnTypeIdent = mScriptingData.mIdentifierMap->GetIdentifier(_returnType, WNStrLen(_returnType));
        nativeReturn = mScriptingData.mTypeManager->GetNativeFromWNType(returnTypeIdent);
    }
    if(!returnTypeIdent) {
        return(WN_INVALID_IDENTIFIER);
    }
    
    if (!nativeReturn) {
        return(WN_INVALID_TYPE);
    }

    std::vector<WNFunction*> previouslyMapped;
    GetExternalFunctions(functionName, previouslyMapped, WN_FALSE);

    std::vector<WNFunction*>::iterator i = previouslyMapped.begin();
    for(; i != previouslyMapped.end(); ++i) {
        if((*i)->mTypes.size() != functionTypes.size()) {
            continue;
        }
        WN_UINT32 param = 0;
        for(; param < (*i)->mTypes.size(); ++param) {
            if((*i)->mTypes[param] != functionTypes[param]) {
                break;
            }
        }
        if (param == (*i)->mTypes.size()) {
            break;
        }
    }

    if(i != previouslyMapped.end()) {
        return(WN_CONFLICTING_DEFINITION);
    }

    WNFunction* func = WN_NEW(WNFunction());

    func->mFunctionName = correctFunctionName;
    func->mCallingConv = WNCallingConvention::WN_CDECL;
    func->mExternal = WN_TRUE;
    func->mFunctionPointer = _funcPtr;
    func->mReturnType = nativeReturn;
    func->mStackRet = _stack;

    for(std::vector<WNNativeElementType*>::iterator i = functionTypes.begin(); i != functionTypes.end(); ++i) {
        func->mTypes.push_back(*i);
    }

    GenerateLLVMFunction(func, WN_TRUE);

    mNativeMap.insert(std::make_pair(functionName, func));
    
    return(WN_OK);
}


WNFunctionManager::ReturnCode WNFunctionManager::AddExternalFunction(WN_VOID* _funcPtr, WNIdentifierElement _name, WNTypeElement _returnType, std::vector<WNTypeElement>& _params) {
    std::vector<WNFunction*> previouslyMapped;
    GetExternalFunctions(_name, previouslyMapped, WN_FALSE);

    std::vector<WNFunction*>::iterator i = previouslyMapped.begin();
    for(; i != previouslyMapped.end(); ++i) {
        if((*i)->mTypes.size() != _params.size()) {
            continue;
        }
        WN_UINT32 param = 0;
        for(; param < (*i)->mTypes.size(); ++param) {
            if((*i)->mTypes[param]->mScriptType != _params[param]) {
                break;
            }
        }
        if (param == (*i)->mTypes.size()) {
            break;
        }
    }

    if(i != previouslyMapped.end()) {
        return(WN_CONFLICTING_DEFINITION);
    }
    
    WNFunction* func = WN_NEW(WNFunction());
    func->mFunctionName = _name;
    func->mCallingConv = WNCallingConvention::WN_CDECL;
    func->mFunctionPointer = _funcPtr;
    func->mStackRet = WN_FALSE;
    for(std::vector<WNTypeElement>::iterator i = _params.begin(); i != _params.end(); ++i) {
        WNNativeElementType* elemType = mScriptingData.mTypeManager->GetNativeFromType(*i);
        if (!elemType) {
            WN_DELETE(func);
            return(WN_INVALID_TYPE);
        }
        func->mTypes.push_back(elemType);
    }
    GenerateLLVMFunction(func, WN_TRUE);
    func->mExternal = WN_TRUE;

    mNativeMap.insert(std::make_pair(func->mFunctionName, func));

    return(WN_OK);
}

WN_VOID WNFunctionManager::GetExternalFunctions(WNIdentifierElement _name, std::vector<WNFunction*>& _functions, WN_BOOL _eTypes) {
    std::pair<FunctionMMap::iterator, FunctionMMap::iterator> its = mNativeMap.equal_range(_name);

    while(its.first != its.second && _name == ((_eTypes)?its.first->first:its.first->second->mFunctionName)){
        _functions.push_back(its.first->second);
        ++its.first;
    }

    return;
}

WN_VOID WNFunctionManager::GenerateLLVMFunction(WNFunction* _func, WN_BOOL _external) {
    std::vector<llvm::Type*> arguments;
    for(std::vector<WNNativeElementType*>::iterator i = _func->mTypes.begin();
        i != _func->mTypes.end(); ++i) {
        arguments.push_back(mScriptingData.mTypeManager->GetLLVMType((*i)->mScriptType));
    }
    llvm::Type* t;
    if(_func->mStackRet) {
        t = mScriptingData.mTypeManager->GetIntPtrType();
    } else {
        t = mScriptingData.mTypeManager->GetLLVMType(_func->mReturnType->mScriptType);
    }
    llvm::FunctionType* ftype = llvm::FunctionType::get(t, arguments, WN_FALSE);
    llvm::Function* f = llvm::Function::Create(ftype, 
                                               (_external?llvm::Function::ExternalLinkage : llvm::Function::InternalLinkage),
                                               mScriptingData.mIdentifierMap->GetStringFromIdentifier(_func->mFunctionName),
                                               mScriptingData.mModule);
    _func->mLlvmFunction = f;
    _func->mLlvmFunctiontype = ftype;
}

WNIdentifierMap* WNFunctionManager::GetIdentifierMap() {
    return(mScriptingData.mIdentifierMap);
}

WNFunctionManager::ReturnCode WNFunctionManager::AddIncompleteScriptFunction(WNIdentifierElement _name, WNIdentifierElement _file, WN_BOOL _onStack, WNTypeElement _returnType, const std::vector<WNTypeElement>& _params, WN_SIZE_T _offset, WN_SIZE_T _endOffset) {
	std::pair<FunctionMMap::iterator, FunctionMMap::iterator> its = mFunctionMap.equal_range(_name);
    std::vector<const WNFunction*> mFunctions;
    while(its.first != its.second && its.first->second->mFunctionName == _name) {
        mFunctions.push_back(its.first->second);
    	++its.first;
	}
    for(std::vector<const WNFunction*>::iterator i = mFunctions.begin(); i != mFunctions.end(); ++i) {
        //We are not the "same" function if we are not defined in the same file or have a difference in 
        //parameters. For now going to stick with the convention of (2 functions are the same if they have the same
        //parameters but different return types)
        // TODO: Do not allow the function to be registered if we have defined this function in a file
        // that this file references
        if ((*i)->mTypes.size() != _params.size() || 
            (*i)->mFile != _file) {
            continue;
        }
        WN_BOOL exists = WN_TRUE;
        for(size_t ct = 0; ct < (*i)->mTypes.size(); ++ct) {
            if ((*i)->mTypes[ct]->mScriptType != _params[ct ]) {
                exists = WN_FALSE;
                break;
            }
        }
        if (exists) {
            return(WNFunctionManager::WN_CONFLICTING_DEFINITION);
        }
    }

    WNFunction* f = WN_NEW(WNFunction());
    f->mFunctionName = _name;
    f->mFile = _file;
    f->mExternal = WN_FALSE;
    f->mCallingConv = WNCallingConvention::WN_CDECL;
    f->mFunctionPointer = WN_NULL;
    f->mLlvmFunction = WN_NULL;
    f->mReturnType = mScriptingData.mTypeManager->GetNativeFromType(_returnType);
    f->mOffset = _offset;
    f->mEndOffset = _endOffset;
    f->mStackRet = _onStack;
    if (!f->mReturnType) {
        WN_DELETE(f);
        return(WN_INVALID_TYPE);
    }
    for(size_t i = 0; i < _params.size(); ++i) {
        WNNativeElementType * nativeType = mScriptingData.mTypeManager->GetNativeFromType(_params[i]);
        if (!nativeType) {
            WN_DELETE(f);
            return(WN_INVALID_TYPE);
        }
        f->mTypes.push_back(nativeType);
    }
    GenerateLLVMFunction(f, WN_FALSE);
    mFunctionMap.insert(std::make_pair(f->mFunctionName, f));
    return(WN_OK);
}


WN_BOOL WNFunctionManager::ContainsFile(IncludeMMap::iterator& _start, IncludeMMap::iterator& _end, WNIdentifierElement _element) {
    for(IncludeMMap::iterator i = _start; i != _end; ++i) {
        if(i->second == _element){ return(true); }
    }
    return(false);
}

WN_VOID WNFunctionManager::GetScriptFunctions(std::vector<WNFunction*>& _functions, WNIdentifierElement _name, WNIdentifierElement _file) {
	std::pair<FunctionMMap::iterator, FunctionMMap::iterator> its = mFunctionMap.equal_range(_name);
    std::pair<IncludeMMap::iterator, IncludeMMap::iterator> files = mIncludeMap.equal_range(_file);
    //Look for functions with the correct name
    
    while(its.first != its.second && 
          its.first->second->mFunctionName == _name &&
          ContainsFile(files.first, files.second, _file)) { //if they are in the set of files we are talking about return
        _functions.push_back(its.first->second);
        ++its.first;
    }

    GetExternalFunctions(_name, _functions, WN_TRUE);
}

WNFunction* WNFunctionManager::GetCreateFunctionFunction() {
    return(mCreateFunctionfunction);
}

WNFunction* WNFunctionManager::GetMallocFunction() {
    return(mMallocFunction);
}

WNFunction* WNFunctionManager::GetFreeFunction() {
    return(mFreeFunction);
}

WNFunction* WNFunctionManager::GetMemcpyFunction() {
    return(mMemcpyFunction);
}

WNFunction* WNFunctionManager::GetTemporaryStackFunction() {
    return(mInternalTempStackFunction);
}

WNFunctionManager::ReturnCode WNFunctionManager::GetExactScriptFunction(WNFunction*& _function, WNIdentifierElement _name, const std::vector<WNTypeElement>& _params, WNIdentifierElement _file) {
    std::vector<WNFunction*> functions;
    GetScriptFunctions(functions, _name, _file);
    if (functions.size() == 0){
        return(WN_DOES_NOT_EXIST);
    }
    _function = WN_NULL;
    for(size_t i = 0; i < functions.size(); ++i){
        WNFunction* f = functions[i];
        if (f->mTypes.size() != _params.size()) {
            continue;
        }
        WN_BOOL identical = WN_TRUE;
        for(size_t k = 0; k < _params.size(); ++k) {
            if (f->mTypes[i]->mScriptType != _params[i]) {
                identical = WN_FALSE;
                break;
            }
        }

        if (!identical) {
            break;
        }
        if (_function) {
            return(WN_IDENTICAL_DUPLICATES);
        }
        _function = f;
    }
    return(WN_OK);
}

WN_VOID WNFunctionManager::CreateInternalFunction(WNFunction*& _funcLocation, const WN_CHAR* _functionName, WN_VOID* _funcPtr, llvm::Type* _ptrType, std::vector<llvm::Type*> _llvmTypes ) {
    _funcLocation = WN_NEW(WNFunction());
    _funcLocation->mExternal = WN_TRUE;
    _funcLocation->mFunctionPointer = _funcPtr;
    llvm::FunctionType* ftype = llvm::FunctionType::get(_ptrType, _llvmTypes, WN_FALSE);
    llvm::Function* f = llvm::Function::Create(ftype, 
        llvm::Function::ExternalLinkage, 
        _functionName,
        mScriptingData.mModule);
    _funcLocation->mLlvmFunction = f;
    _funcLocation->mCallingConv = WNCallingConvention::WN_CDECL;
    _funcLocation->mLlvmFunctiontype = ftype;
    _funcLocation->mStackRet = WN_FALSE;
}

WN_VOID WNFunctionManager::AddNewThread() {
    tlFunctionManager = this;
    tlTempStackLocation = WN_NULL;
    tlTempStackSize = 0;
}
